//
//  UICollectionView+YHExtension.swift
//  Adorbee
//
//  Created by amovision on 2025/5/13.
//

import UIKit

public extension UICollectionView {
  var numOfSections: (() -> (Int))? {
    set { objc_setAssociatedObject(self, #function.pointer, newValue, .OBJC_ASSOCIATION_RETAIN_NONATOMIC) }
    get { objc_getAssociatedObject(self, #function.pointer) as? () -> (Int) }
  }

  var numOfRows: ((Int) -> (Int))? {
    set { objc_setAssociatedObject(self, #function.pointer, newValue, .OBJC_ASSOCIATION_RETAIN_NONATOMIC) }
    get { objc_getAssociatedObject(self, #function.pointer) as? (Int) -> (Int) }
  }

  var cellForRows: ((Int, Int) -> UICollectionViewCell)? {
    set { objc_setAssociatedObject(self, #function.pointer, newValue, .OBJC_ASSOCIATION_RETAIN_NONATOMIC) }
    get { objc_getAssociatedObject(self, #function.pointer) as? (Int, Int) -> UICollectionViewCell }
  }

  var didSelectRow: ((Int, Int) -> Void)? {
    set { objc_setAssociatedObject(self, #function.pointer, newValue, .OBJC_ASSOCIATION_RETAIN_NONATOMIC) }
    get { objc_getAssociatedObject(self, #function.pointer) as? (Int, Int) -> Void }
  }

  var sizeForRows: ((Int, Int) -> CGSize)? {
    set { objc_setAssociatedObject(self, #function.pointer, newValue, .OBJC_ASSOCIATION_RETAIN_NONATOMIC) }
    get { objc_getAssociatedObject(self, #function.pointer) as? (Int, Int) -> CGSize }
  }

  var sizeForHeaders: ((UICollectionViewLayout, Int) -> CGSize)? {
    set { objc_setAssociatedObject(self, #function.pointer, newValue, .OBJC_ASSOCIATION_RETAIN_NONATOMIC) }
    get { objc_getAssociatedObject(self, #function.pointer) as? (UICollectionViewLayout, Int) -> CGSize }
  }

  var sizeForFooters: ((UICollectionViewLayout, Int) -> CGSize)? {
    set { objc_setAssociatedObject(self, #function.pointer, newValue, .OBJC_ASSOCIATION_RETAIN_NONATOMIC) }
    get { objc_getAssociatedObject(self, #function.pointer) as? (UICollectionViewLayout, Int) -> CGSize }
  }

  var sizeForHeader: CGSize {
    set { objc_setAssociatedObject(self, #function.pointer, newValue, .OBJC_ASSOCIATION_RETAIN_NONATOMIC) }
    get { objc_getAssociatedObject(self, #function.pointer) as? CGSize ?? .zero }
  }

  var sizeForFooter: CGSize {
    set { objc_setAssociatedObject(self, #function.pointer, newValue, .OBJC_ASSOCIATION_RETAIN_NONATOMIC) }
    get { objc_getAssociatedObject(self, #function.pointer) as? CGSize ?? .zero }
  }

  var viewForHeaders: ((Int) -> UIView)? {
    set { objc_setAssociatedObject(self, #function.pointer, newValue, .OBJC_ASSOCIATION_RETAIN_NONATOMIC) }
    get { objc_getAssociatedObject(self, #function.pointer) as? (Int) -> UIView }
  }

  var viewForFooters: ((Int) -> UIView)? {
    set { objc_setAssociatedObject(self, #function.pointer, newValue, .OBJC_ASSOCIATION_RETAIN_NONATOMIC) }
    get { objc_getAssociatedObject(self, #function.pointer) as? (Int) -> UIView }
  }

  private(set) var scrollDirection: UICollectionView.ScrollDirection {
    set { objc_setAssociatedObject(self, #function.pointer, newValue, .OBJC_ASSOCIATION_ASSIGN) }
    get { objc_getAssociatedObject(self, #function.pointer) as? UICollectionView.ScrollDirection ?? .vertical }
  }

  var flowLayout: UICollectionViewFlowLayout? {
    collectionViewLayout as? UICollectionViewFlowLayout
  }

  var registerCells: [UICollectionViewCell.Type]? {
    set { objc_setAssociatedObject(self, #function.pointer, newValue, .OBJC_ASSOCIATION_RETAIN_NONATOMIC) }
    get { objc_getAssociatedObject(self, #function.pointer) as? [UICollectionViewCell.Type] }
  }

  convenience init(cells: UICollectionViewCell.Type..., layout: UICollectionViewLayout, scrollDirection: UICollectionView.ScrollDirection = .vertical) {
    self.init(frame: .zero, collectionViewLayout: layout)
    backgroundColor = .clear
    showsVerticalScrollIndicator = false
    showsHorizontalScrollIndicator = false
    self.scrollDirection = scrollDirection
    registerCells = cells.map { $0 }
    cells.forEach { register($0, forCellWithReuseIdentifier: NSStringFromClass($0)) }
    register(headers: YHCollectionReusableView.self, footers: YHCollectionReusableView.self)
    delegate = self
    dataSource = self
  }

  convenience init(cells: UICollectionViewCell.Type..., itemSize: CGSize, rowColumn: CGSize = .zero, sectionInsets: UIEdgeInsets = .zero, scrollDirection: UICollectionView.ScrollDirection = .vertical) {
    let layout = UICollectionViewFlowLayout()
    layout.sectionHeadersPinToVisibleBounds = true
    layout.itemSize = itemSize
    layout.scrollDirection = scrollDirection
    layout.minimumLineSpacing = rowColumn.height
    layout.minimumInteritemSpacing = rowColumn.width
    layout.sectionInset = sectionInsets
    self.init(frame: .zero, collectionViewLayout: layout)
    self.scrollDirection = scrollDirection
    backgroundColor = .clear
    showsVerticalScrollIndicator = false
    showsHorizontalScrollIndicator = false
    registerCells = cells.map { $0 }
    cells.forEach { register($0, forCellWithReuseIdentifier: NSStringFromClass($0)) }
    register(headers: YHCollectionReusableView.self, footers: YHCollectionReusableView.self)
    delegate = self
    dataSource = self
  }
}

extension UICollectionView: @retroactive UIScrollViewDelegate {}
extension UICollectionView: @retroactive UICollectionViewDelegate, @retroactive UICollectionViewDataSource, @retroactive UICollectionViewDelegateFlowLayout {
  public func collectionView(_ collectionView: UICollectionView, numberOfItemsInSection section: Int) -> Int {
    numOfRows?(section) ?? 1
  }

  public func numberOfSections(in collectionView: UICollectionView) -> Int {
    numOfSections?() ?? 0
  }

  public func collectionView(_ collectionView: UICollectionView, cellForItemAt indexPath: IndexPath) -> UICollectionViewCell {
    cellForRows?(indexPath.section, indexPath.item) ?? UICollectionViewCell()
  }

  public func collectionView(_ collectionView: UICollectionView, didSelectItemAt indexPath: IndexPath) {
    didSelectRow?(indexPath.section, indexPath.row)
  }

  public func collectionView(_ collectionView: UICollectionView, layout collectionViewLayout: UICollectionViewLayout, sizeForItemAt indexPath: IndexPath) -> CGSize {
    sizeForRows?(indexPath.section, indexPath.item) ?? flowLayout?.itemSize ?? .zero
  }

  public func collectionView(_ collectionView: UICollectionView, viewForSupplementaryElementOfKind kind: String, at indexPath: IndexPath) -> UICollectionReusableView {
    if kind == UICollectionView.elementKindSectionHeader {
      let header = collectionView.dequeueReusableHeader(named: YHCollectionReusableView.self, for: indexPath)
      header.contentView = viewForHeaders?(indexPath.section) ?? UIView(frame: CGRect(origin: .zero, size: sizeForHeader))
      return header
    } else {
      let footer = collectionView.dequeueReusableFooter(named: YHCollectionReusableView.self, for: indexPath)
      footer.contentView = viewForFooters?(indexPath.section) ?? UIView(frame: CGRect(origin: .zero, size: sizeForFooter))
      return footer
    }
  }

  public func collectionView(_ collectionView: UICollectionView, layout collectionViewLayout: UICollectionViewLayout, referenceSizeForHeaderInSection section: Int) -> CGSize {
    sizeForHeaders?(collectionViewLayout, section) ?? sizeForHeader
  }

  public func collectionView(_ collectionView: UICollectionView, layout collectionViewLayout: UICollectionViewLayout, referenceSizeForFooterInSection section: Int) -> CGSize {
    sizeForFooters?(collectionViewLayout, section) ?? sizeForFooter
  }
}

public extension UICollectionView {
  func register<T: UICollectionViewCell>(cells: T.Type...) {
    for cell in cells {
      register(cell, forCellWithReuseIdentifier: NSStringFromClass(cell))
    }
  }

  func register<T: UICollectionReusableView>(headers: T.Type..., footers: T.Type...) {
    for header in headers {
      register(header, forSupplementaryViewOfKind: UICollectionView.elementKindSectionHeader, withReuseIdentifier: NSStringFromClass(header))
    }
    for footer in footers {
      register(footer, forSupplementaryViewOfKind: UICollectionView.elementKindSectionFooter, withReuseIdentifier: NSStringFromClass(footer))
    }
  }

  func dequeueReusableCell<T: UICollectionViewCell>(named name: T.Type, for indexPath: IndexPath, configCell: ((T) -> Void)? = nil) -> T {
    guard let cell = dequeueReusableCell(withReuseIdentifier: NSStringFromClass(name), for: indexPath) as? T else {
      fatalError(
        "Couldn't find UICollectionViewCell for \(NSStringFromClass(name)), make sure the cell is registered with collection view")
    }
    configCell?(cell)
    return cell
  }

  func dequeueReusableCell<T: UICollectionViewCell>(named name: T.Type, section: Int, item: Int, configCell: ((T) -> Void)? = nil) -> T {
    guard let cell = dequeueReusableCell(withReuseIdentifier: NSStringFromClass(name), for: IndexPath(item: item, section: section)) as? T else {
      fatalError(
        "Couldn't find UICollectionViewCell for \(NSStringFromClass(name)), make sure the cell is registered with collection view")
    }
    configCell?(cell)
    return cell
  }

  func dequeueReusableHeader<T: UICollectionReusableView>(named name: T.Type, for indexPath: IndexPath) -> T {
    guard let view = dequeueReusableSupplementaryView(ofKind: UICollectionView.elementKindSectionHeader, withReuseIdentifier: NSStringFromClass(name), for: indexPath) as? T else {
      fatalError(
        "Couldn't find UICollectionReusableView for \(NSStringFromClass(name)), make sure the view is registered with collection view")
    }
    return view
  }

  func dequeueReusableFooter<T: UICollectionReusableView>(named name: T.Type, for indexPath: IndexPath) -> T {
    guard let view = dequeueReusableSupplementaryView(ofKind: UICollectionView.elementKindSectionFooter, withReuseIdentifier: NSStringFromClass(name), for: indexPath) as? T else {
      fatalError(
        "Couldn't find UICollectionReusableView for \(NSStringFromClass(name)), make sure the view is registered with collection view")
    }
    return view
  }
}
